package com.javen.zookeeper.lock;

import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;

/**
 * @author Javen
 */
@Slf4j
public class DistributedLock implements Lock, Watcher {
    private ZooKeeper zk;
    /**
     * 根
     */
    private final String root = "/locks";
    /**
     * 竞争资源的标志
     */
    private final String lockName;
    /**
     * 等待前一个锁的节点
     */
    private String waitNode;
    /**
     * 当前锁的节点
     */
    private String currentNode;
    /**
     * 计数器
     */
    private CountDownLatch latch;
    private final int sessionTimeout = 30000;
    private final List<Exception> exception = new ArrayList<>();

    

    /**
     * 创建分布式锁
     *
     * @param connectString 127.0.0.1:2181,127.0.0.1:2182
     * @param lockName      竞争资源标志，lockName 中不能包含单词 lock
     */
    public DistributedLock(String connectString, String lockName) {
        this.lockName = lockName;
        // 创建一个与服务器的连接
        try {
            zk = new ZooKeeper(connectString, sessionTimeout, this);
            Stat stat = zk.exists(root, false);
            if (stat == null) {
                // 创建根节点
                zk.create(root, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
        } catch (IOException | KeeperException | InterruptedException e) {
            exception.add(e);
        }
    }

    /**
     * zookeeper节点的监视器
     */
    @Override
    public void process(WatchedEvent event) {
        if (event.getState() == Event.KeeperState.Closed && this.latch != null) {
            this.latch.countDown();
        }
    }

    @Override
    public void lock() {
        if (exception.size() > 0) {
            throw new LockException(exception.get(0));
        }
        try {
            if (this.tryLock()) {
                log.info("Thread " + Thread.currentThread().getId() + " " + currentNode + " get lock true");
            } else {
                // 等待锁
                waitForLock(waitNode, sessionTimeout);
            }
        } catch (KeeperException | InterruptedException e) {
            throw new LockException(e);
        }
    }

    @Override
    public boolean tryLock() {
        try {
            String splitStr = "_lock_";
            if (lockName.contains(splitStr)) {
                throw new LockException("lockName can not contains lock");
            }
            // 创建临时子节点
            String separator = "/";
            currentNode = zk.create(root + separator + lockName + splitStr,
                    new byte[0],
                    ZooDefs.Ids.OPEN_ACL_UNSAFE,
                    CreateMode.EPHEMERAL_SEQUENTIAL);
            log.info(currentNode + " is created ");
            // 取出所有子节点
            List<String> subNodes = zk.getChildren(root, false);
            // 取出所有lockName的锁
            List<String> lockObjNodes = new ArrayList<>();
            for (String node : subNodes) {
                String tempNode = node.split(splitStr)[0];
                if (tempNode.equals(lockName)) {
                    lockObjNodes.add(node);
                }
            }
            Collections.sort(lockObjNodes);
            if (currentNode.equals(root + separator + lockObjNodes.get(0))) {
                // 如果是最小的节点,则表示取得锁
                return true;
            }
            // 如果不是最小的节点，找到比自己小1的节点
            String subCurrentNode = currentNode.substring(currentNode.lastIndexOf(separator) + 1);
            waitNode = lockObjNodes.get(Collections.binarySearch(lockObjNodes, subCurrentNode) - 1);
        } catch (KeeperException | InterruptedException e) {
            throw new LockException(e);
        }
        return false;
    }

    @Override
    public boolean tryLock(long time, TimeUnit unit) {
        try {
            if (this.tryLock()) {
                return true;
            }
            return waitForLock(waitNode, time);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }

    private boolean waitForLock(String lower, long waitTime) throws InterruptedException, KeeperException {
        Stat stat = zk.exists(root + "/" + lower, true);
        // 判断比自己小一个数的节点是否存在,如果不存在则无需等待锁,同时注册监听
        if (stat != null) {
            log.info("Thread " + Thread.currentThread().getId() + " waiting for " + root + "/" + lower);
            this.latch = new CountDownLatch(1);
            this.latch.await(waitTime, TimeUnit.MILLISECONDS);
            this.latch = null;
        }
        return true;
    }

    @Override
    public void unlock() {
        try {
            log.info("unlock " + currentNode);
            zk.delete(currentNode, -1);
            currentNode = null;
            zk.close();
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
        }
    }

    @Override
    public Condition newCondition() {
        return null;
    }

    @Override
    public void lockInterruptibly() {
        this.lock();
    }

    public static class LockException extends RuntimeException {
        private static final long serialVersionUID = 1L;

        public LockException(String e) {
            super(e);
        }

        public LockException(Exception e) {
            super(e);
        }
    }
}